#include<iostream>
#include<random>
#include<time.h>
using namespace std;

/*
extern "C" int func_name(int nparam, void **params, int *ndims, int64_t **shapes, const char **dtypes, void *stream, void *extra);
nparam (int): 输入输出总数。比如算子有2个输入，1个输出，则nparam的值为3。
params (void **): 输入输出指针数组。比如算子有2个输入，1个输出，params[0]指向第一个输入数据，params[1]指向第二个输入数据，params[2]指向输出数据
ndims (int *): 输入输出shape维度数组。比如params[i]是个shape[1024, 1024]的张量，则ndims[i]的值为2。
shapes (int64_t **): 输入输出shape数组。比如params[i]是个shape[1024, 1024]的张量，则shapes[i][0]的值为1024，shapes[i][1]的值为1024。
dtypes (const char **): 输入输出数据类型数组。dtypes里的元素取值可为：“float32”, “float16”, “float”, “float64”, “int”, “int8”, “int16”,
 “int32”, “int64”, “uint”, “uint8”, “uint16”, “uint32”, “uint64”, “bool”。
stream (void *): CUDA流指针，仅定义GPU算子实现时需要。
extra (void *): 用于后续扩展。
*/
extern "C" int CustomSigmoidRandomDrop(int nparam, void **params, int *ndims, 
int64_t **shapes, const char **dtypes, void *stream, void *extra){
    //nparam=5
    //nidms:[4,4,4,1,1]
    //shapes:shape[0][0-3]=[1,1,712,1072]
    //shapes:shape[1][0-3]=[1,1,712,1072]
    //shapes:shape[2][0-1]=[1,1,712,1072]
    //shapes:shape[3][0]=[2]
    //shapes:shape[4][0]=[1]
    if (nparam!=5) return 1;
    float *data=static_cast<float *>(params[0]);
    float *label=static_cast<float *>(params[1]);
    int *indicator=static_cast<int *>(params[2]);
    int *num_area=static_cast<int *>(params[3]);
    float *out=static_cast<float *>(params[4]);
    
    int length=shapes[0][2]*shapes[0][3];
    //int length=712*1072;
    double p_forground=0.0;//损失函数前一项概率总和
    double p_background=0.0;//损失函数后一项概率总和
    int num_background_reserved=0;//保留的背景点数量
    int num_forground=0;//前景区元素数量

    mt19937 gen((unsigned int) time(nullptr));
    uniform_real_distribution<float>dis(0,1);

    for(int i=0;i<length;i++){
        double prob=(double)data[i];//data[i]即为判别概率
        if(label[i]==0){
            //当前像素的位置为标签中背景点像素位置
            //根据dis(gen)的随机值决定是否保留像素点
            if(dis(gen)<1-prob) continue;
            else{
                indicator[i]=2;//保留的背景点在判别矩阵中值为2
                num_background_reserved++;
                p_background+=log((double)(1-prob));       
            }
        }
        else{
            num_forground++;
            indicator[i]=1;//前景点在判别矩阵中值为1
            p_forground+=log(prob);
        }
    }
    num_area[0]=num_forground;
    num_area[1]=num_background_reserved;
    double beta=num_background_reserved/(num_background_reserved+num_forground);
    double result=-beta*p_forground-(1-beta)*p_background;
    out[0]=(float)result;
    return 0;//仅作为判断函数是否执行正确
}